import requests
import time
import re

MAILDROP_GRAPHQL_URL = "https://api.maildrop.cc/graphql"
HEADERS = {"content-type": "application/json"}

def graphql_query(query, variables=None):
    body = {"query": query}
    if variables:
        body["variables"] = variables
    resp = requests.post(MAILDROP_GRAPHQL_URL, headers=HEADERS, json=body)
    resp.raise_for_status()
    return resp.json()

def get_inbox(mailbox):
    query = """
    query GetInbox($mailbox: String!) {
        inbox(mailbox: $mailbox) {
            id
            headerfrom
            subject
            date
        }
    }
    """
    return graphql_query(query, {"mailbox": mailbox})["data"]["inbox"]

def get_message(mailbox, msg_id):
    query = """
    query GetMessage($mailbox: String!, $id: String!) {
        message(mailbox: $mailbox, id: $id) {
            id
            subject
            data
            html
        }
    }
    """
    return graphql_query(query, {"mailbox": mailbox, "id": msg_id})["data"]["message"]

def extract_code(text):
    # 根据实际验证码规则调整正则
    match = re.search(r"\b([A-Za-z0-9]{4,8})\b", text)
    return match.group(1) if match else None

def poll_for_code(mailbox, subject_keyword="验证码", timeout=180, poll_interval=5):
    print(f"轮询邮箱 {mailbox}@maildrop.cc，等待包含“{subject_keyword}”的新邮件...")
    start = time.time()
    seen_ids = set()
    while time.time() - start < timeout:
        inbox = get_inbox(mailbox)
        for mail in inbox:
            if mail["id"] in seen_ids:
                continue
            if subject_keyword in mail.get("subject", ""):
                print(f"收到目标邮件: {mail['subject']}")
                msg = get_message(mailbox, mail["id"])
                print("邮件正文:\n", msg["data"])
                code = extract_code(msg["data"] or "")
                print("提取到的验证码:", code)
                return code
            seen_ids.add(mail["id"])
        print("未收到目标邮件，继续等待...")
        time.sleep(poll_interval)
    print("超时未收到目标邮件")
    return None

if __name__ == "__main__":
    mailbox = "your_mailbox"  # 只填前缀，不带@maildrop.cc
    poll_for_code(mailbox, subject_keyword="验证码")